package workerpool

// workerpool 的实现：
// 1. pool 的创建和销毁
// 2. pool 中 worker（goroutine）的调度
// 3. task 的提交与调度

import (
	"errors"
	"fmt"
	"sync"
)

const (
	defaultCapacity int = 10
	maxCapacity     int = 20
)

type Task func()

type Pool struct {
	capacity int // workerpool 的大小，即 worker 的容量

	// 带缓冲的 channel，作为 worker 的计数器
	// 当 active 可写时，就创建一个 worker，用于处理用户通过 Schedule 函数提交的待处理 Task
	active chan struct{}

	// Schedule 函数将 Task 提交到 tasks，已创建的 worker 会从这里读取 task 并执行
	tasks chan Task

	// 用于在 pool 销毁时，等待所有 worker 退出
	wg sync.WaitGroup

	// 用于通知各个 worker 退出的信号 channel
	quit chan struct{}
}

func New(capacity int) *Pool {
	if capacity <= 0 {
		capacity = defaultCapacity
	}

	if capacity > maxCapacity {
		capacity = maxCapacity
	}

	p := &Pool{
		capacity: capacity,
		tasks:    make(chan Task),
		quit:     make(chan struct{}),

		active: make(chan struct{}, capacity),
	}

	fmt.Println("worker pool started")

	go p.run()

	return p
}

// 哨兵错误
var ErrWorkerPoolFreed = errors.New("worker pool freed")

// 将要执行的任务添加到 workerpool 的 task channel 中
func (p *Pool) Schedule(t Task) error {
	select {
	case <-p.quit:
		return ErrWorkerPoolFreed
	case p.tasks <- t:
		return nil
	}
}

func (p *Pool) Free() {
	close(p.quit)
	p.wg.Wait()

	fmt.Println("workerpool freed")
}

func (p *Pool) run() {
	idx := 0

	for {
		select {
		case <-p.quit: // 当 quit channel 收到信号，该 goroutine 就会退出，意味着 workerpool 不会再创建新 worker
			return
		case p.active <- struct{}{}:
			// create a new worker, idx 作为 worker 的编号
			idx++
			p.newWorker(idx)
		}
	}
}

func (p *Pool) newWorker(i int) {
	// 每次创建一个 newWorker (goroutine)， sync.WaitGroup 递增
	p.wg.Add(1)

	go func() {
		defer func() {
			if err := recover(); err != nil {
				fmt.Printf("worker[%03d]: recover panic[%s] and exit\n", i, err)
				<-p.active // 若 worker 由于 panic 退出，则 active 计数器会减1，然后 run 方法里收到信号，会再创建一个新的 worker
			}
			p.wg.Done()
		}()

		fmt.Printf("worker[%03d] start\n", i)

		for {
			select {
			case <-p.quit:
				fmt.Printf("worker[%03d] exit\n", i)
				<-p.active
				return
			case t := <-p.tasks:
				fmt.Printf("worker[%03d]: received a task\n", i)
				t()
			}
		}
	}()
}
